import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import sys
import os
import netCDF4 as nc
import numpy as np

# 处理打包后的资源路径
if getattr(sys, 'frozen', False):
    base_path = sys._MEIPASS
    cartopy.config['data_dir'] = os.path.join(base_path, 'cartopy', 'data')

# 确保中文显示
plt.rcParams["font.family"] = ["SimSun"]


def plot_geo_heatmap(file_path, var_name, lat_var_name=None, lon_var_name=None):
    """
    绘制指定HDF4文件中特定变量的热力图

    参数:
    file_path: str - 数据文件的路径
    var_name: str - 要绘制的变量名
    lat_var_name: str, optional - 纬度变量名（手动指定）
    lon_var_name: str, optional - 经度变量名（手动指定）
    """
    try:
        # 打开文件
        print(f"正在打开HDF4文件: {file_path}")
        ds = nc.Dataset(file_path, 'r')
        print(f"成功打开文件。文件中的所有变量: {list(ds.variables.keys())}")

        # 检查变量是否存在
        if var_name not in ds.variables:
            ds.close()
            raise ValueError(f"变量 {var_name} 不存在于文件中。可用变量: {list(ds.variables.keys())}")

        # 读取变量数据
        var_data = ds.variables[var_name][:]
        print(f"变量 {var_name} 的形状: {var_data.shape}")

        # 尝试获取经纬度变量
        lat_data = None
        lon_data = None

        # 确定经纬度变量名
        if lat_var_name and lon_var_name:
            if lat_var_name in ds.variables and lon_var_name in ds.variables:
                lat_data = ds.variables[lat_var_name][:]
                lon_data = ds.variables[lon_var_name][:]
                print(f"成功获取经纬度变量: {lat_var_name}, {lon_var_name}")
            else:
                print(f"无法获取指定的经纬度变量: {lat_var_name}, {lon_var_name}")
        else:
            # 自动查找经纬度变量
            lat_vars = ["lat", "Latitude", "Lat", "y", "Y"]
            lon_vars = ["lon", "Longitude", "Lon", "x", "X"]
            print(f"正在自动查找经纬度变量: 纬度变量列表={lat_vars}, 经度变量列表={lon_vars}")

            for lat_var in lat_vars:
                if lat_var in ds.variables:
                    lat_data = ds.variables[lat_var][:]
                    print(f"成功获取纬度变量: {lat_var}")
                    break

            for lon_var in lon_vars:
                if lon_var in ds.variables:
                    lon_data = ds.variables[lon_var][:]
                    print(f"成功获取经度变量: {lon_var}")
                    break

        ds.close()

        # 判断变量类型
        is_geo2d = False
        var_type = None

        if lat_data is not None and lon_data is not None:
            is_geo2d = True
            var_type = "Geo2D"
            print("识别为地理数据")
        elif len(var_data.shape) == 1:
            var_type = "1D"
        elif len(var_data.shape) == 2:
            var_type = "2D"
            print("识别为普通2D数据")
        else:
            # 处理高维变量
            # 这里简化处理，直接取第一个二维切片
            if len(var_data.shape) > 2:
                # 假设前两个维度是空间维度
                var_data = var_data[0, :, :] if var_data.shape[0] < var_data.shape[1] else var_data[:, 0, :]
                print(f"已选择前两个维度进行绘图，新形状: {var_data.shape}")
                var_type = "2D"
            else:
                raise ValueError(f"变量 {var_name} 维度过多且无法自动选择绘图维度")

        print(f"最终变量类型: {var_type}")

        # 清除 matplotlib 中的旧图形
        plt.clf()

        # 创建新图表
        fig = plt.figure(figsize=(10, 6))

        # 根据变量类型绘图
        if var_type == "1D":
            plt.plot(var_data)
            plt.title(f"{var_name} - 折线图")
        elif var_type in ["2D", "Geo2D"]:
            if var_type == "Geo2D" and lat_data is not None and lon_data is not None:
                # 确保经纬度数据形状与变量数据匹配
                if lat_data.ndim == 1 and lon_data.ndim == 1:
                    lon_grid, lat_grid = np.meshgrid(lon_data, lat_data)
                else:
                    lon_grid = lon_data
                    lat_grid = lat_data

                ax = fig.add_subplot(111, projection=ccrs.PlateCarree())
                im = ax.pcolormesh(
                    lon_grid,
                    lat_grid,
                    var_data,
                    transform=ccrs.PlateCarree(),
                    cmap="viridis"
                )
                ax.coastlines()
                ax.gridlines(draw_labels=True)
                plt.colorbar(im, label="值")
                plt.title(f"{var_name} - 地理热力图")
            else:
                plt.imshow(var_data, cmap="viridis")
                plt.colorbar(label="值")
                plt.title(f"{var_name} - 热力图")
        else:
            raise ValueError(f"变量 {var_name} 类型不支持绘图 ({var_type})")

        plt.tight_layout()
        plt.show()

    except Exception as e:
        error_msg = f"绘图错误: {str(e)}"
        print(error_msg)
        raise


def main():
    """主函数，用于演示如何调用plot_geo_heatmap函数"""
    # 在这里设置文件路径和变量名
    file_path = "G:\MOD11C1\MOD11C1.A2004001.061.2020168035507.hdf"
    var_name = "LST_Day_CMG"

    # 手动指定经纬度变量名
    plot_geo_heatmap(file_path, var_name, lat_var_name="Lat", lon_var_name="Lon")


if __name__ == "__main__":
    main()